/*
*    usage - Print impact on system, resources used, of specified user.
*    Copyright (C) 2024  Marcus Pedersén marcus@marcux.org
*
*    This program is free software: you can redistribute it and/or modify
*    it under the terms of the GNU General Public License as published by
*    the Free Software Foundation, either version 3 of the License, or
*    (at your option) any later version.
*
*    This program is distributed in the hope that it will be useful,
*    but WITHOUT ANY WARRANTY; without even the implied warranty of
*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
*    GNU General Public License for more details.
*
*    You should have received a copy of the GNU General Public License
*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/


package main

import (
	"fmt"
	"os"
	"os/user"
	"unicode"
	"bufio"
	"strings"
	"strconv"
	"github.com/spf13/cobra"
)

const VERSION string = "1.0.0"
const AUTHORS string = "Marcus Pedersen <marcus@marcux.org>"
const AUTHNAME string = "Marcus Pedersen"


type Process struct {
	uid string
	pid string
	ramKb int
	ussRam int
	swapKb int
	threads int
}


// Returns true if
// string only contains digits
func StrOnlyDigits(s string) bool {
	for _, c := range s {
		if !unicode.IsDigit(c) {
			return false
		}
	}

	return true
}


// Prints help text and exit
func PrintVersion() {
	fmt.Println("usage ", VERSION)
	fmt.Println("Copyright (C) 2024 Marcus Pedersen")
	fmt.Println("License GPLv3+: GNU GPL version 3 or later <https://gnu.org/licenses/gpl.html>")
	fmt.Println("This is free software: you are free to change and redistribute it.")
	fmt.Println("There is NO WARRANTY, to the extent permitted by law.")
	fmt.Println("")
	fmt.Println("Written by ", AUTHNAME)
}


func main() {
	var uid string
	var psList []Process = make([]Process, 0)

	var cName string
	var cUid string
	var cVersion bool
	var cHelp bool

	rootCmd := &cobra.Command {
		Use: "usage [OPTION]...",
		Short: "Print impact on system, resources used, of specified user.",
		Long: `Prints usage of resources
for a single user. If no user is
specified with flags the current
user usage is printed.
RSS - Resident Set Size
all memory used inclusive shared memory
and pages and other common memory.
USS - Unique Set Size
represent the private memory of a process.`,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) > 0 {
				fmt.Println("Positional arguments not available")
				fmt.Println("Try -h or --help for help.")
				os.Exit(3)
			}
		},
	}

    rootCmd.Flags().StringVarP(&cName, "name", "n", "", "`Name` of user to show info about")
	rootCmd.Flags().StringVarP(&cUid, "uid", "u", "", "`Uid` of user to show info about")
	rootCmd.Flags().BoolVarP(&cVersion, "version", "V", false, "Show version text and exit")
	rootCmd.Flags().BoolVarP(&cHelp, "help", "h", false, "Show this help text and exit")

    rootCmd.SilenceErrors = true
    rootCmd.SilenceUsage = true
    if _, err := rootCmd.ExecuteC(); err != nil {
        fmt.Println("Argument error:", err)
        fmt.Println("Try -h or --help for help")
        os.Exit(1)
    }

	if cHelp {
		os.Exit(0)
	}

	if cVersion {
		PrintVersion()
		os.Exit(0)
	}

	if cName != "" && cUid != "" {
		fmt.Println("Argument error: --name and --uid can not be used together.")
		fmt.Println("Try -h or --help for help.")
		os.Exit(4)
	}
	
	var usr *user.User

	if cName != "" {
		luser, err := user.Lookup(cName)
		if err != nil {
			fmt.Printf("Failed to lookup user: %s\n%s\n",
				cName, err)
			os.Exit(5)
		}

		usr = luser
	} else if cUid != "" {
		iuser, err := user.LookupId(cUid)
		if err != nil {
			fmt.Printf("Failed to lookup user uid: %s\n%s\n",
				cUid, err)
			os.Exit(6)
		}

		usr = iuser
	} else {
		currUser, err := user.Current()
		if err != nil {
			fmt.Println("Failed to get current user")
			fmt.Println(err)
			os.Exit(7)
		}
	
		usr = currUser
	}

	uid = usr.Uid
	fmt.Printf("Username:       %20s\n", usr.Username)
	fmt.Printf("UID:                      %10s\n", usr.Uid)


	dirEnt, err := os.ReadDir("/proc")
	if err != nil {
		fmt.Printf("Failed to open directory: %s", err)
	}

	for _, ent := range dirEnt {
		if ent.IsDir() {
			if StrOnlyDigits(ent.Name()) {
				f, err := os.Open("/proc/" + ent.Name() + "/status")
				if err != nil {
					fmt.Println("Failed to open file : ", "/proc/"+ent.Name())
				}
				defer f.Close()

				scan := bufio.NewScanner(f)

				p := new(Process)
				isUserPs := false
				for scan.Scan() {
					if strings.Contains(scan.Text(), "Uid:") {
						fieldUid := strings.Fields(scan.Text())
						if len(fieldUid) >= 2 {
							if fieldUid[1] == uid {
								isUserPs = true
								p.uid = fieldUid[1]
								p.pid = ent.Name()
							}
						}
					}
					if strings.Contains(scan.Text(), "Threads:") {
						fieldThreads := strings.Fields(scan.Text())
						if len(fieldThreads) == 2 {
							if t, err := strconv.Atoi(fieldThreads[1]); err == nil {
								p.threads = t
							}
						}
					}

					if strings.Contains(scan.Text(), "VmRSS:") {
						fieldVmSize := strings.Fields(scan.Text())
						if len(fieldVmSize) == 3 {
							if r, err := strconv.Atoi(fieldVmSize[1]); err == nil {
								p.ramKb = r
							}
						}
					}

					if strings.Contains(scan.Text(), "VmSwap:") {
						fieldVmSwap := strings.Fields(scan.Text())
						if len(fieldVmSwap) == 3 {
							if s, err := strconv.Atoi(fieldVmSwap[1]); err == nil {
								p.swapKb = s
							}
						}
					}
				}
				if isUserPs {
					fd, _ := os.Open("/proc/" + ent.Name() + "/smaps_rollup")
					//if err != nil {
					//	fmt.Println("Failed to open file : ", "/proc/"+ent.Name()+"/snaps_rollup")
					//}
					defer fd.Close()

					scan := bufio.NewScanner(fd)

					for scan.Scan() {
						if strings.HasPrefix(scan.Text(), "Private") {
							list_priv := strings.Fields(scan.Text())
							if len(list_priv) == 3 {
								if pr, err := strconv.Atoi(list_priv[1]); err == nil {
									p.ussRam += pr
								}
							}
						}
					}

					psList = append(psList, *p)
				}
			}
		}
	}

	noProc := len(psList)
	ram, swap, threads, uss := 0, 0, 0, 0
	for _, p := range psList {
		ram += p.ramKb
		uss += p.ussRam
		swap += p.swapKb
		threads += p.threads
	}

	fmt.Println("")
	fmt.Printf("Total number of processes:%10d\n", noProc)
	fmt.Printf("Total number of threads:  %10d\n", threads)
	fmt.Printf("Total RAM (RSS) usage:    %10d kB\n", ram)
	fmt.Printf("Total RAM (USS) usage:    %10d kB\n", uss)
	fmt.Printf("Total swap usage:         %10d kB\n", swap)
}
